/*
    PEB Offset calculator
    @5mukx
*/


use ntapi::{
    ntpebteb::PEB,
    ntrtl::RTL_USER_PROCESS_PARAMETERS,
};
use std::arch::asm;

use memoffset::offset_of;


fn get_peb() -> *mut PEB {
    #[cfg(target_arch = "x86_64")]
    {
        let peb: *mut PEB;
        unsafe {asm!("mov {}, gs:[0x60]", out(reg) peb)};
        peb
    }
    #[cfg(target_arch = "x86")]
    {
        let peb: *mut PEB;
        unsafe { asm!("mov {}, fs:[0x30]", out(reg) peb)};
        peb
    }
}  

fn print_peb_offsets() {
    println!("[*] PEB Field Offsets:");
    println!("----Field: Offset----");
    println!("[+] InheritedAddressSpace: {:#x}", offset_of!(PEB, InheritedAddressSpace));
    println!("[+] ReadImageFileExecOptions: {:#x}", offset_of!(PEB, ReadImageFileExecOptions));
    println!("[+] BeingDebugged: {:#x}", offset_of!(PEB, BeingDebugged));
    println!("[+] BitField: {:#x}", offset_of!(PEB, BitField));
    println!("[+] Mutant: {:#x}", offset_of!(PEB, Mutant));
    println!("[+] ImageBaseAddress: {:#x}", offset_of!(PEB, ImageBaseAddress));
    println!("[+] Ldr: {:#x}", offset_of!(PEB, Ldr));
    println!("[+] ProcessParameters: {:#x}", offset_of!(PEB, ProcessParameters));
    println!("[+] SubSystemData: {:#x}", offset_of!(PEB, SubSystemData));
    println!("[+] ProcessHeap: {:#x}", offset_of!(PEB, ProcessHeap));
    println!("[+] FastPebLock: {:#x}", offset_of!(PEB, FastPebLock));
    println!("[+] AtlThunkSListPtr: {:#x}", offset_of!(PEB, AtlThunkSListPtr));
    println!("[+] IFEOKey: {:#x}", offset_of!(PEB, IFEOKey));
    println!("[+] CrossProcessFlags: {:#x}", offset_of!(PEB, CrossProcessFlags));
    
    // u handle filed
    let u_offset = offset_of!(PEB, u);
    println!("[+] u.KernelCallbackTable: {:#x}", u_offset);
    println!("[+] u.UserSharedInfoPtr: {:#x}", u_offset);
    println!("[+] SystemReserved: {:#x}", offset_of!(PEB, SystemReserved));
    println!("[+] AtlThunkSListPtr32: {:#x}", offset_of!(PEB, AtlThunkSListPtr32));
    println!("[+] ApiSetMap: {:#x}", offset_of!(PEB, ApiSetMap));
    println!("[+] TlsExpansionCounter: {:#x}", offset_of!(PEB, TlsExpansionCounter));
    println!("[+] TlsBitmap: {:#x}", offset_of!(PEB, TlsBitmap));
    println!("[+] TlsBitmapBits: {:#x}", offset_of!(PEB, TlsBitmapBits));
    println!("[+] ReadOnlySharedMemoryBase: {:#x}", offset_of!(PEB, ReadOnlySharedMemoryBase));
    println!("[+] SharedData: {:#x}", offset_of!(PEB, SharedData));
    println!("[+] ReadOnlyStaticServerData: {:#x}", offset_of!(PEB, ReadOnlyStaticServerData));
    println!("[+] AnsiCodePageData: {:#x}", offset_of!(PEB, AnsiCodePageData));
    println!("[+] OemCodePageData: {:#x}", offset_of!(PEB, OemCodePageData));
    println!("[+] UnicodeCaseTableData: {:#x}", offset_of!(PEB, UnicodeCaseTableData));
    println!("[+] NumberOfProcessors: {:#x}", offset_of!(PEB, NumberOfProcessors));
    println!("[+] NtGlobalFlag: {:#x}", offset_of!(PEB, NtGlobalFlag));
    println!("[+] CriticalSectionTimeout: {:#x}", offset_of!(PEB, CriticalSectionTimeout));
    println!("[+] HeapSegmentReserve: {:#x}", offset_of!(PEB, HeapSegmentReserve));
    println!("[+] HeapSegmentCommit: {:#x}", offset_of!(PEB, HeapSegmentCommit));
    println!("[+] HeapDeCommitTotalFreeThreshold: {:#x}", offset_of!(PEB, HeapDeCommitTotalFreeThreshold));
    println!("[+] HeapDeCommitFreeBlockThreshold: {:#x}", offset_of!(PEB, HeapDeCommitFreeBlockThreshold));
    println!("[+] NumberOfHeaps: {:#x}", offset_of!(PEB, NumberOfHeaps));
    println!("[+] MaximumNumberOfHeaps: {:#x}", offset_of!(PEB, MaximumNumberOfHeaps));
    println!("[+] ProcessHeaps: {:#x}", offset_of!(PEB, ProcessHeaps));
    println!("[+] GdiSharedHandleTable: {:#x}", offset_of!(PEB, GdiSharedHandleTable));
    println!("[+] ProcessStarterHelper: {:#x}", offset_of!(PEB, ProcessStarterHelper));
    println!("[+] GdiDCAttributeList: {:#x}", offset_of!(PEB, GdiDCAttributeList));
    println!("[+] LoaderLock: {:#x}", offset_of!(PEB, LoaderLock));
    println!("[+] OSMajorVersion: {:#x}", offset_of!(PEB, OSMajorVersion));
    println!("[+] OSMinorVersion: {:#x}", offset_of!(PEB, OSMinorVersion));
    println!("[+] OSBuildNumber: {:#x}", offset_of!(PEB, OSBuildNumber));
    println!("[+] OSCSDVersion: {:#x}", offset_of!(PEB, OSCSDVersion));
    println!("[+] OSPlatformId: {:#x}", offset_of!(PEB, OSPlatformId));
    println!("[+] ImageSubsystem: {:#x}", offset_of!(PEB, ImageSubsystem));
    println!("[+] ImageSubsystemMajorVersion: {:#x}", offset_of!(PEB, ImageSubsystemMajorVersion));
    println!("[+] ImageSubsystemMinorVersion: {:#x}", offset_of!(PEB, ImageSubsystemMinorVersion));
    println!("[+] ActiveProcessAffinityMask: {:#x}", offset_of!(PEB, ActiveProcessAffinityMask));
    println!("[+] GdiHandleBuffer: {:#x}", offset_of!(PEB, GdiHandleBuffer));
    println!("[+] PostProcessInitRoutine: {:#x}", offset_of!(PEB, PostProcessInitRoutine));
    println!("[+] TlsExpansionBitmap: {:#x}", offset_of!(PEB, TlsExpansionBitmap));
    println!("[+] TlsExpansionBitmapBits: {:#x}", offset_of!(PEB, TlsExpansionBitmapBits));
    println!("[+] SessionId: {:#x}", offset_of!(PEB, SessionId));
    println!("[+] AppCompatFlags: {:#x}", offset_of!(PEB, AppCompatFlags));
    println!("[+] AppCompatFlagsUser: {:#x}", offset_of!(PEB, AppCompatFlagsUser));
    println!("[+] pShimData: {:#x}", offset_of!(PEB, pShimData));
    println!("[+] AppCompatInfo: {:#x}", offset_of!(PEB, AppCompatInfo));
    println!("[+] CSDVersion: {:#x}", offset_of!(PEB, CSDVersion));
    println!("[+] ActivationContextData: {:#x}", offset_of!(PEB, ActivationContextData));
    println!("[+] ProcessAssemblyStorageMap: {:#x}", offset_of!(PEB, ProcessAssemblyStorageMap));
    println!("[+] SystemDefaultActivationContextData: {:#x}", offset_of!(PEB, SystemDefaultActivationContextData));
    println!("[+] SystemAssemblyStorageMap: {:#x}", offset_of!(PEB, SystemAssemblyStorageMap));
    println!("[+] MinimumStackCommit: {:#x}", offset_of!(PEB, MinimumStackCommit));
    println!("[+] FlsCallback: {:#x}", offset_of!(PEB, FlsCallback));
    println!("[+] FlsListHead: {:#x}", offset_of!(PEB, FlsListHead));
    println!("[+] FlsBitmap: {:#x}", offset_of!(PEB, FlsBitmap));
    println!("[+] FlsBitmapBits: {:#x}", offset_of!(PEB, FlsBitmapBits));
    println!("[+] FlsHighIndex: {:#x}", offset_of!(PEB, FlsHighIndex));
    println!("[+] WerRegistrationData: {:#x}", offset_of!(PEB, WerRegistrationData));
    println!("[+] WerShipAssertPtr: {:#x}", offset_of!(PEB, WerShipAssertPtr));
    println!("[+] pUnused: {:#x}", offset_of!(PEB, pUnused));
    println!("[+] pImageHeaderHash: {:#x}", offset_of!(PEB, pImageHeaderHash));
    println!("[+] TracingFlags: {:#x}", offset_of!(PEB, TracingFlags));
    println!("[+] CsrServerReadOnlySharedMemoryBase: {:#x}", offset_of!(PEB, CsrServerReadOnlySharedMemoryBase));
    println!("[+] TppWorkerpListLock: {:#x}", offset_of!(PEB, TppWorkerpListLock));
    println!("[+] TppWorkerpList: {:#x}", offset_of!(PEB, TppWorkerpList));
    println!("[+] WaitOnAddressHashTable: {:#x}", offset_of!(PEB, WaitOnAddressHashTable));
    println!("[+] TelemetryCoverageHeader: {:#x}", offset_of!(PEB, TelemetryCoverageHeader));
    println!("[+] CloudFileFlags: {:#x}", offset_of!(PEB, CloudFileFlags));
    println!("[+] CloudFileDiagFlags: {:#x}", offset_of!(PEB, CloudFileDiagFlags));
    println!("[+] PlaceholderCompatibilityMode: {:#x}", offset_of!(PEB, PlaceholderCompatibilityMode));
    println!("[+] PlaceholderCompatibilityModeReserved: {:#x}", offset_of!(PEB, PlaceholderCompatibilityModeReserved));
    println!("[+] LeapSecondData: {:#x}", offset_of!(PEB, LeapSecondData));
    println!("[+] LeapSecondFlags: {:#x}", offset_of!(PEB, LeapSecondFlags));
    println!("[+] NtGlobalFlag2: {:#x}", offset_of!(PEB, NtGlobalFlag2));
}

fn print_rtl_user_process_parameters_offsets() {
    println!("[*] RTL_USER_PROCESS_PARAMETERS Field Offsets:");
    println!("----Field: Offset----");
    println!("[+] MaximumLength: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, MaximumLength));
    println!("[+] Length: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, Length));
    println!("[+] Flags: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, Flags));
    println!("[+] DebugFlags: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, DebugFlags));
    println!("[+] ConsoleHandle: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, ConsoleHandle));
    println!("[+] ConsoleFlags: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, ConsoleFlags));
    println!("[+] StandardInput: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, StandardInput));
    println!("[+] StandardOutput: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, StandardOutput));
    println!("[+] StandardError: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, StandardError));
    println!("[+] CurrentDirectory: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, CurrentDirectory));
    println!("[+] DllPath: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, DllPath));
    println!("[+] ImagePathName: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, ImagePathName));
    println!("[+] CommandLine: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, CommandLine));
    println!("[+] Environment: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, Environment));
    println!("[+] StartingX: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, StartingX));
    println!("[+] StartingY: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, StartingY));
    println!("[+] CountX: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, CountX));
    println!("[+] CountY: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, CountY));
    println!("[+] CountCharsX: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, CountCharsX));
    println!("[+] CountCharsY: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, CountCharsY));
    println!("[+] FillAttribute: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, FillAttribute));
    println!("[+] WindowFlags: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, WindowFlags));
    println!("[+] ShowWindowFlags: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, ShowWindowFlags));
    println!("[+] WindowTitle: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, WindowTitle));
    println!("[+] DesktopInfo: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, DesktopInfo));
    println!("[+] ShellInfo: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, ShellInfo));
    println!("[+] RuntimeData: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, RuntimeData));
    println!("[+] CurrentDirectories: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, CurrentDirectories));
    println!("[+] EnvironmentSize: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, EnvironmentSize));
    println!("[+] EnvironmentVersion: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, EnvironmentVersion));
    println!("[+] PackageDependencyData: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, PackageDependencyData));
    println!("[+] ProcessGroupId: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, ProcessGroupId));
    println!("[+] LoaderThreads: {:#x}", offset_of!(RTL_USER_PROCESS_PARAMETERS, LoaderThreads));
}


// example runtime
fn validate_offsets_runtime() {
    unsafe{
        println!("\n[*] Runtime Offset Validation:");
        let peb = get_peb();
        if peb.is_null() {
            println!("[*] PEB is null");
            return;
        }
        let peb_base = peb as usize;
        let process_parameters = (*peb).ProcessParameters;
        if process_parameters.is_null() {
            println!("[*] ProcessParameters is null");
            return;
        }
        let params_base = process_parameters as usize;

        // PEB.ProcessParameters
        let process_parameters_offset = (process_parameters as usize).wrapping_sub(peb_base);
        println!("[*] PEB.ProcessParameters: Calculated Offset = {:#x}, Expected = {:#x}",
                process_parameters_offset, offset_of!(PEB, ProcessParameters));

        // PEB.u.KernelCallbackTable
        let kernel_callback_table_offset = (&(*peb).u as *const _ as usize).wrapping_sub(peb_base);
        println!("[*] PEB.u.KernelCallbackTable: Calculated Offset = {:#x}, Expected = {:#x}",
                kernel_callback_table_offset, offset_of!(PEB, u));
        
        // Example !
        // RTL_USER_PROCESS_PARAMETERS fields
        let current_directory_offset = (&(*process_parameters).CurrentDirectory as *const _ as usize).wrapping_sub(params_base);
        let command_line_offset = (&(*process_parameters).CommandLine as *const _ as usize).wrapping_sub(params_base);
        let environment_offset = (&(*process_parameters).Environment as *const _ as usize).wrapping_sub(params_base);

        println!("[*] RTL_USER_PROCESS_PARAMETERS.CurrentDirectory: Calculated Offset = {:#x}, Expected = {:#x}",
                current_directory_offset, offset_of!(RTL_USER_PROCESS_PARAMETERS, CurrentDirectory));

        println!("[*] RTL_USER_PROCESS_PARAMETERS.CommandLine: Calculated Offset = {:#x}, Expected = {:#x}",
                command_line_offset, offset_of!(RTL_USER_PROCESS_PARAMETERS, CommandLine));

        println!("[*] RTL_USER_PROCESS_PARAMETERS.Environment: Calculated Offset = {:#x}, Expected = {:#x}",
                environment_offset, offset_of!(RTL_USER_PROCESS_PARAMETERS, Environment));

        println!("[*] PEB Address: {:#x}", peb_base);
        println!("[*] ProcessParameters Address: {:#x}", process_parameters as usize);
        println!("[*] u.KernelCallbackTable Address: {:#x}", (*peb).u.KernelCallbackTable as usize);
        println!("[*] Environment Address: {:#x}", (*process_parameters).Environment as usize);
    }
}


fn main(){
    
    print_peb_offsets();
    
    print_rtl_user_process_parameters_offsets();
    
    validate_offsets_runtime();
    
}